import json
import time
import requests
import logging
from typing import Dict, Union

# 豆包模型API配置
DOUBAO_API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
DOUBAO_MODEL = "doubao-seed-1-6-flash-250615"

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("dspy_doubao.log")
    ]
)
logger = logging.getLogger("DoubaoModel")


class DoubaoLM:
    """豆包模型集成类"""

    def __init__(
        self,
        api_key: str,
        model: str = DOUBAO_MODEL,
        max_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        request_timeout: int = 30,
        max_history: int = 10,
    ):
        # 配置模型参数
        self.api_key = api_key
        self.model_name = model
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.request_timeout = request_timeout
        self.max_history = max_history

        # 初始化状态
        self.history = []
        self.total_requests = 0
        self.failed_requests = 0
        self.last_request_time = 0
        self.request_rate = 5

        # 配置请求头
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        logger.info(f"初始化豆包模型: {model}")

    def _rate_limit(self):
        """实现请求速率限制"""
        current_time = time.time()
        elapsed = current_time - self.last_request_time

        min_interval = 1.0 / self.request_rate
        if elapsed < min_interval:
            sleep_time = min_interval - elapsed
            logger.debug(f"速率限制: 等待{sleep_time:.2f}秒")
            time.sleep(sleep_time)

        self.last_request_time = time.time()

    def generate_response(self, prompt: str) -> str:
        """生成响应 - 直接调用豆包API"""
        # 添加用户消息到历史
        self.history.append({"role": "user", "content": prompt})

        # 准备API请求数据
        payload = {
            "model": self.model_name,
            "messages": self.history[-self.max_history:],
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p
        }

        # 重试机制
        for attempt in range(self.max_retries + 1):
            self._rate_limit()
            logger.debug(f"发送请求到豆包API (尝试 {attempt+1}/{self.max_retries+1})")

            response = None
            try:
                response = requests.post(
                    DOUBAO_API_URL,
                    headers=self.headers,
                    data=json.dumps(payload),
                    timeout=self.request_timeout
                )

                response.raise_for_status()
                result = response.json()

            except requests.exceptions.HTTPError as e:
                self.failed_requests += 1
                status_code = response.status_code if response else -1
                if status_code == 401:
                    logger.error("认证失败: 无效的API密钥")
                    return "认证错误: 请检查API密钥"
                elif status_code == 429:
                    logger.warning("请求过多: 达到速率限制")
                    backoff = self.retry_delay * (2 ** attempt)
                    time.sleep(backoff)
                    continue
                else:
                    logger.error(f"HTTP错误 {status_code}: {str(e)}")
                    if response:
                        logger.error(f"响应内容: {response.text}")
                    return "API请求失败"

            except Exception as e:
                self.failed_requests += 1
                logger.error(f"请求异常: {str(e)}")
                return "API请求失败"

            if attempt == self.max_retries:
                break

            # 解析响应
            try:
                if "choices" in result and len(result["choices"]) > 0:
                    ai_response: str = result["choices"][0]["message"]["content"].strip()
                elif "output" in result:
                    ai_response: str = result["output"].strip()
                else:
                    raise ValueError(f"无法解析API响应: {json.dumps(result, indent=2)}")
            except Exception as e:
                logger.error(f"解析异常: {str(e)}")
                return "解析AI响应失败"

            # 添加AI响应到历史
            self.history.append({"role": "assistant", "content": ai_response})
            self.total_requests += 1

            logger.debug(f"API响应成功: {ai_response[:50]}...")
            return ai_response

        logger.error(f"所有重试均失败: {prompt}")
        return "抱歉，当前无法处理您的请求"

    def get_stats(self) -> Dict[str, Union[int, float]]:
        """获取模型使用统计"""
        success_rate = 100.0
        if self.total_requests > 0:
            success_rate = (self.total_requests - self.failed_requests) / self.total_requests * 100

        return {
            "total_requests": self.total_requests,
            "failed_requests": self.failed_requests,
            "success_rate": success_rate
        }

    def reset_history(self):
        """重置对话历史"""
        self.history = []
        logger.info("对话历史已重置")
